import torch
import torch.nn as nn
from torch.autograd import Variable

device = "cuda" if torch.cuda.is_available else "cpu"

class ConvLSTMCell(nn.Module):
    def __init__(self, input_channels, hidden_channels, kernel_size):
        super(ConvLSTMCell, self).__init__()

        assert hidden_channels % 2 == 0

        self.input_channels = input_channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.num_features = 4

        self.padding = int((kernel_size - 1) / 2)

        self.Wxi = nn.Conv2d(self.input_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=True)
        self.Whi = nn.Conv2d(self.hidden_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=False)
        self.Wxf = nn.Conv2d(self.input_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=True)
        self.Whf = nn.Conv2d(self.hidden_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=False)
        self.Wxc = nn.Conv2d(self.input_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=True)
        self.Whc = nn.Conv2d(self.hidden_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=False)
        self.Wxo = nn.Conv2d(self.input_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=True)
        self.Who = nn.Conv2d(self.hidden_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=False)

        self.Wci = None
        self.Wcf = None
        self.Wco = None

    def forward(self, x, h, c):  # x: 上层Cnn输出的conv1
        ci = torch.sigmoid(self.Wxi(x) + self.Whi(h) + c * self.Wci)  # [5, 32, 30, 30]
        cf = torch.sigmoid(self.Wxf(x) + self.Whf(h) + c * self.Wcf)
        cc = cf * c + ci * torch.tanh(self.Wxc(x) + self.Whc(h))
        co = torch.sigmoid(self.Wxo(x) + self.Who(h) + cc * self.Wco)
        ch = co * torch.tanh(cc)
        return ch, cc  # H, cell

    def init_hidden(self, batch_size, hidden, shape):
        if self.Wci is None:
            self.Wci = Variable(torch.zeros(1, hidden, shape[0], shape[1])).to("cpu")  # 定义为全0变量; [1, 32, 30, 30]
            self.Wcf = Variable(torch.zeros(1, hidden, shape[0], shape[1])).to("cpu")
            self.Wco = Variable(torch.zeros(1, hidden, shape[0], shape[1])).to("cpu")
        else:
            assert shape[0] == self.Wci.size()[2], 'Input Height Mismatched!'
            assert shape[1] == self.Wci.size()[3], 'Input Width Mismatched!'
        return (Variable(torch.zeros(batch_size, hidden, shape[0], shape[1])).to("cpu"),  # [5, 32, 30, 30]
                Variable(torch.zeros(batch_size, hidden, shape[0], shape[1])).to("cpu"))  # [5, 32, 30, 30]


class ConvLSTM(nn.Module):
    # input_channels corresponds to the first input feature map
    # hidden state is a list of succeeding lstm layers.
    def __init__(self, input_channels, hidden_channels, kernel_size, step=1, effective_step=[1]):
        super(ConvLSTM, self).__init__()
        self.input_channels = [input_channels] + hidden_channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size  # 3
        self.num_layers = len(hidden_channels)  # LSTM层数
        self.step = step
        self.effective_step = effective_step
        self._all_layers = []
        for i in range(self.num_layers):  # n层lstm
            name = 'cell{}'.format(i)
            cell = ConvLSTMCell(self.input_channels[i], self.hidden_channels[i], self.kernel_size)  # 每层lstm
            setattr(self, name, cell)
            self._all_layers.append(cell)

    def forward(self, input):  # [5, 32, 30, 30]; 上层conv输出的conv1
        internal_state = []
        outputs = []
        for step in range(self.step):  # lstm5个时间步
            x = input[step:(step+1)]
            for i in range(self.num_layers):  # 遍历多层lstm
                # all cells are initialized in the first step
                name = 'cell{}'.format(i)
                if step == 0:  # 第一个时间步，没有h和c，要先生成h和c
                    bsize, _, height, width = x.size()  # 如果是初始的lstm，定义初始的h和c值
                    (h, c) = getattr(self, name).init_hidden(batch_size=bsize, hidden=self.hidden_channels[i],    # getattr(self, name):ConvLSTMCell
                                                             shape=(height, width))
                    internal_state.append((h, c))

                # do forward
                (h, c) = internal_state[i]  # 提取h和c
                x, new_c = getattr(self, name)(x, h, c)  # ConvLSTMCell(x, h, c)
                internal_state[i] = (x, new_c)  # 赋值到internal_state; 供下个步骤使用
                outputs.append(x[0])

        return [torch.stack(outputs,dim=0)], (x, new_c)


        # 下面为原始代码，个人认为没有使用LSTM的性质，所以做了上面的更改(如果实际问题中，修改后代码存在问题，可用源码)
        # internal_state = []
        # outputs = []
        # for step in range(self.step):  # lstm5个时间步
        #     x = input
        #     for i in range(self.num_layers):  # 遍历多层lstm
        #         # all cells are initialized in the first step
        #         name = 'cell{}'.format(i)
        #         if step == 0:  # 第一个时间步，没有h和c，要先生成h和c
        #             bsize, _, height, width = x.size()  # 如果是初始的lstm，定义初始的h和c值
        #             (h, c) = getattr(self, name).init_hidden(batch_size=bsize, hidden=self.hidden_channels[i],
        #                                                      # getattr(self, name):ConvLSTMCell
        #                                                      shape=(height, width))
        #             internal_state.append((h, c))
        #
        #         # do forward
        #         (h, c) = internal_state[i]  # 提取h和c
        #         x, new_c = getattr(self, name)(x, h, c)  # ConvLSTMCell(x, h, c)
        #         internal_state[i] = (x, new_c)  # 赋值到internal_state; 供下个步骤使用
        #     # only record effective steps
        #     if step in self.effective_step:  # 保留最后一层的lstm结果
        #         outputs.append(x)
        #
        # return outputs, (x, new_c)


if __name__ == '__main__':
    # gradient check
    convlstm = ConvLSTM(input_channels=512, hidden_channels=[128, 64, 64, 32, 32], kernel_size=3, step=5,
                        effective_step=[4]).to("cpu")
    loss_fn = torch.nn.MSELoss()

    input = Variable(torch.randn(5, 512, 64, 32)).to("cpu")
    target = Variable(torch.randn(1, 32, 64, 32)).double().to("cpu")

    output = convlstm(input)
    output = output[0][0].double()
    print(output.shape)
    # res = torch.autograd.gradcheck(loss_fn, (output, target), eps=1e-6, raise_exception=True)
    # print(res)
